#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@File        : mngs_cleanup.py
@Author      : Bing Liang
@Email       : believer19940901@gmail.com
@Date        : 2025/9/28
@Description : 宏基因组数据磁盘空间清理脚本（支持 dry-run / execute 模式）
"""

import argparse
import time
import subprocess
from pathlib import Path
from multiprocessing import Pool, cpu_count

# ===== 配置区 =====
# 需要清理的目标目录（分析结果）
BASE_PATH = Path("/data/mNGS/runmngs/result")

# 原始数据目录
RAW_DATA_PATH = Path("/data/mNGS/runmngs/raw_fq")

# 备份目录
BACK_PATH = Path("/mnt/mngs/users/mNGS/rawFq_ZD")

# 一个月时间（秒）
MONTH = 30 * 24 * 60 * 60 * 2

# 清理目标的子文件（相对于 run/sample 目录）
SAMPLE_CLEAN_TARGETS = [
    "{sample}.sorted.bam",
    "{sample}.sorted.bam.csi",
    "{sample}.unmap.fastq",
    "{sample}.removepla.fasta",
    "{sample}.removedup.raw.fasta",
    "{sample}.txt",
    "{sample}_merged.sam",
    "{sample}.uniq_stat_tmp",
    "{sample}_plas_and_homo.tmp",
    "{sample}.merge_litesam_tmp",
    "{sample}_hit_and_genus.tmp"
]


# ===== 辅助函数 =====
def build_cleanup_commands():
    """生成清理命令列表"""
    cmds = []

    # 1️⃣ 遍历分析结果目录
    for run_dir in BASE_PATH.glob("*"):
        if not run_dir.is_dir():
            continue

        age = time.time() - run_dir.stat().st_mtime
        if age <= MONTH:
            continue  # 只清理一个月前的目录

        run_id = run_dir.stem

        # 删除大文件与报告
        cmds.append(f"rm -f '{run_dir / 'all.final.fasta'}'")
        cmds.append(f"rm -f '{run_dir / f'{run_id}.report.tar.gz'}'")
        cmds.append(f"rm -rf '{run_dir / run_id / 'reports'}'")

        # 删除样本中间文件
        for sub_dir in run_dir.iterdir():

            # 不删除fastq目录（因为里面只是软连接，不占空间）
            if sub_dir.stem == "fastq":
                continue

            # 删除shell目录（shell脚本和运行日志）,不用保留
            if sub_dir.stem == "shell":
                cmds.append(f"rm -rf '{sub_dir}'")
                continue

            # 不不删除结果文件
            if not sub_dir.is_dir() or sub_dir.stem == run_id:
                continue

            sample_id = sub_dir.stem
            for pattern in SAMPLE_CLEAN_TARGETS:
                file_path = sub_dir / pattern.format(sample=sample_id)
                cmds.append(f"rm -f '{file_path}'")

            # 删除一些中间文件
            cmds.append(f"rm -f '{sub_dir}/*.sam_tmp'")
            cmds.append(f"rm -f '{sub_dir}/*.tmp'")

    # 2️⃣ 遍历原始数据目录
    for data_dir in RAW_DATA_PATH.glob("*"):

        # 跳过不是目录的路径
        if not data_dir.is_dir():
            continue

        # 只删除一个月以前的数据
        if (time.time() - data_dir.stat().st_mtime) <= MONTH:
            continue

        # 获取芯片号
        run_id = data_dir.stem

        # 在备份目录中找这个芯片号
        backup_matches = list(BACK_PATH.rglob(run_id))
        # 如果找到了就删除原始目录中芯片号目录
        if backup_matches:
            cmds.append(f"rm -rf '{data_dir}'")

    return cmds


def write_cleanup_script(commands, script_path="./cleanup.sh"):
    """将清理命令写入 shell 脚本，方便审阅"""
    with open(script_path, "w", encoding="utf-8") as f:
        f.write("#!/bin/bash\n\nset -euo pipefail\n\n")
        for cmd in commands:
            f.write(f"{cmd}\n")
    print(f"[INFO] Cleanup script generated at: {script_path}")


def run_command(cmd):
    """在子进程中执行命令"""
    result = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    if result.returncode != 0:
        print(f"[ERROR] Command failed: {cmd}\n{result.stderr}")
    else:
        print(f"[OK] {cmd}")


def execute_cleanup(commands):
    """多进程并行执行清理命令"""
    with Pool(8) as pool:
        pool.map(run_command, commands)


def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(
        description="宏基因组数据空间清理工具\n支持 dry-run 与 execute 模式。",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument(
        "--execute",
        action="store_true",
        default=False,
        help="实际执行清理（默认只是 dry-run 生成 cleanup.sh）"
    )
    return parser.parse_args()


# ===== 主函数 =====
def main():
    args = parse_args()

    print("[INFO] Building cleanup command list...")
    commands = build_cleanup_commands()
    print(f"[INFO] Total cleanup commands: {len(commands)}")

    # 生成 shell 脚本以供检查
    write_cleanup_script(commands)

    if args.execute:
        print("[INFO] Starting cleanup...")
        execute_cleanup(commands)
        print("[INFO] Cleanup finished.")
    else:
        print("[DRY-RUN] No actual deletion performed. Review cleanup.sh before running.")
        print("[HINT] To actually delete, run: python3 mngs_cleanup.py --execute")


if __name__ == "__main__":
    main()
